#!/usr/bin/env python
# encoding: utf-8
"""
gawker_img_fetch

A script to fetch images from a Gawker server instance
"""
PROGRAM = 'gawker_img_fetch'
VERSION = '$Revision$'[11:-2] # Subversion/Mercurial revision
LICENSE = '''
Copyright (c) 2006,2008 W. A. Carrel. All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation 
  and/or other materials provided with the distribution.
* The names of its contributors may not be used to endorse or promote products 
  derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY ITS CONTRIBUTORS ``AS IS'' AND ANY EXPRESS OR
IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

import sys
from optparse import OptionParser
import time
import getopt
import socket
import select
import tempfile
import logging
import logging.handlers
import os

import Zeroconf

POLL_TIMEOUT = 60000 # ms
POLL_NOCONN_TIMEOUT = 5000 # ms
READ_SIZE = 8192
LOG_FORMAT = '%(pathname)s [%(process)d]: %(message)s'

class GawkerCameraReader(object):
  def __init__(self, host_port_info, file_path, dir_path):
    self.host_port_info = host_port_info
    self.path = file_path
    self.dir_path = dir_path
    self.read_buffer = ''
    self.sock = None
    if 'poll' in dir(select):
      # Congrats you have a reasonable python build.
      self.poller = select.poll()
    else:
      # Hmm, no poll, we'll have to do something more universal.
      self.poller = None
  
  def replace_connection(self):
    logging.debug("Replacing connection")
    if self.sock:
      logging.warning("Closing link.")
      self.sock.close()
      if self.poller:
        try:
          self.poller.unregister(self.sock)
        except:
          pass
    self.read_buffer = ''
    host_port = self.host_port_info.GetHostPort()
    if host_port[0] == None or host_port[1] == None:
      logging.warning("No host and port for service.")
      return

    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try: 
      logging.info("Connecting to %s:%d" % (host_port[0],host_port[1]))
      self.sock.connect(host_port)
      logging.debug("Connected.")
    except Exception, e:
      logging.warning('Error connecting to %s:%d : %s' % 
                      (host_port[0],host_port[1], e))
      return
    if self.poller:
      self.poller.register(self.sock, select.POLLIN)
  
  def read_loop(self):
    while 1:
      if not self.sock: # No socket connected
        logging.debug("No socket connected yet")
        if self.host_port_info.ServiceMoved():
          logging.debug("Service moved, replacing connection")
          self.replace_connection()
          logging.debug("Finished connection replacement")
        if self.poller:
          self.poller.poll(POLL_NOCONN_TIMEOUT)
        elif self.sock:
          logging.debug("Started waiting on socket in select")
          select.select([self.sock.fileno()], [], [], POLL_NOCONN_TIMEOUT/1000)
          logging.debug("Finished waiting on socket in select")
        else:
          logging.debug("Started waiting on nothing in select")
          select.select([], [], [], POLL_NOCONN_TIMEOUT/1000)
          logging.debug("Finished waiting on nothing in select")
        continue
      if self.poller:
        ready = self.poller.poll(POLL_TIMEOUT)
        ready_dict = {}
        for item in ready: ready_dict[item[0]]=item[1]
        if len(ready_dict) == 0:
          # We timed out...
          logging.debug("read_loop timeout")
          self.replace_connection()
        elif ready_dict[self.sock.fileno()] & select.POLLIN:
          new_data = self.sock.recv(READ_SIZE)
          if len(new_data) == 0:
            # We got disconnected.
            self.replace_connection()
          else:
            self.read_buffer += new_data
            while self.ParseReadBuffer(): # Loop through the parser
              pass
        else:
          logging.warning("Got unexpected poll() list: %s" % ready)
      else: # No poller for us
        fd = self.sock.fileno()
        logging.debug("Started waiting on socket in select")
        ready = select.select([fd], [], [fd], POLL_TIMEOUT/1000)
        logging.debug("Finished waiting on socket in select")
        if fd in ready[0]:
          new_data = self.sock.recv(READ_SIZE)
          if len(new_data) == 0:
            # We got disconnected.
            self.replace_connection()
          else:
            self.read_buffer += new_data
            while self.ParseReadBuffer():
              pass
        else: # We timed out or had an error/exceptional condition
          logging.debug("error in select loop, reconnecting...")
          self.replace_connection()
  
  def ParseReadBuffer(self):
    if self.read_buffer[:3] == 'DSC': # Description
      block_end = self.read_buffer.find('\n')
      if block_end < 0: 
        # We haven't gotten the full line yet.
        return 0
      self.ParseDescription(self.read_buffer[:block_end])      
      self.read_buffer = self.read_buffer[block_end+1:]
      return 1
    
    if self.read_buffer[:3] == 'IMG': # Image
      line_end = self.read_buffer.find('\n')
      if line_end < 0:
        return 0
      size_start = self.read_buffer.find(':',0,line_end) + 1
      size_end = self.read_buffer.find(':',size_start,line_end)
      image_start = size_end + 2 # Skip the ':' and the '\n'
      size = int(self.read_buffer[size_start:size_end])
      if len(self.read_buffer) < image_start + size:
        # The whole image hasn't arrived yet
        return 0
      self.ReadImage(self.read_buffer[image_start:image_start + size])
      self.read_buffer = self.read_buffer[image_start+size:]
      return 1
      
  def ParseDescription(self, desc_buffer):
    logging.debug("Got description: %s" % (desc_buffer,))
  
  def ReadImage(self, image):
    try:
      (fd, tmppath) = tempfile.mkstemp('.jpg','.gawker_cam_reader', 
        os.path.dirname(self.path))
      if self.dir_path:
        filename = "%04d-%02d-%02d-%02d%02d%02d.jpg" % time.localtime()[:6]
        os.link(tmppath, os.path.join(self.dir_path, filename))
    except OSError, e:
      # If we can't open the temp file we'll throw this image away.
      logging.warning('Unable to create image file: %s' % e)
      return
    try:
      writefrom = 0
      while writefrom < len(image):
        writefrom += os.write(fd,image[writefrom:])
      os.rename(tmppath, self.path)
      logging.debug("Got an image")
      os.close(fd)
    except OSError, e:
      logging.warning('Error writing image file: %s' % e)
      # Clean up after ourselves 
      os.close(fd)
      os.remove(tmppath)


class GawkerListener(object):
  def __init__(self, name=None, host=None, port=None, zeroconf=None):
    self.host = host
    self.port = port
    self.zeroconf = zeroconf
    self.name_to_match = name + '._lapse._tcp.local.'
    self.change_since_last_asked = True
    
    if self.name_to_match:
      logging.debug("Watching zeroconf for '%s'" % (self.name_to_match,))

  def removeService(self, z, t, name):
    if name == self.name_to_match:
      self.host = None
      self.port = None
      self.changed_since_last_asked = True
      logging.debug("Source disappeared")

  def addService(self, zeroconf, type, name):
    self.zeroconf = zeroconf
    if name == self.name_to_match:
      logging.debug("Saw service '%s'" % name)
      self.ProcessInfo()
    else:
      logging.debug("Zeroconf saw '%s', didn't match '%s'" % 
        (name, self.name_to_match))
  
  def ProcessInfo(self):
    if not self.zeroconf: return
    info = self.zeroconf.getServiceInfo('_lapse._tcp.local.', 
      self.name_to_match, 15000)
    if info:
      newhost = str(socket.inet_ntoa(info.getAddress()))
      newport = info.getPort()
      if self.host != newhost or self.port != newport:
        self.host = newhost
        self.port = newport
        self.change_since_last_asked = True
        logging.debug("Source moved to %s:%d" % (self.host, self.port))
    else:
      logging.debug("Unable to obtain service's host and port information.")
    
  
  def ServiceMoved(self):
    self.ProcessInfo() # Check just to make sure...
    return self.change_since_last_asked
  
  def GetHostPort(self):
    self.change_since_last_asked = False
    return (self.host, self.port)

def daemonize():
  # Open the logging goodies early -- just in case
  syslog_handler = logging.handlers.SysLogHandler()
  syslog_handler.setFormatter(logging.Formatter(LOG_FORMAT))
  logging.getLogger('').addHandler(syslog_handler)

  # Borrowed from http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/66012
  # do the UNIX double-fork magic, see Stevens' "Advanced 
  # Programming in the UNIX Environment" for details (ISBN 0201563177)
  try: 
      pid = os.fork() 
      if pid > 0:
          # exit first parent
          sys.exit(0) 
  except OSError, e: 
      print >>sys.stderr, "fork #1 failed: %d (%s)" % (e.errno, e.strerror) 
      sys.exit(1)

  # decouple from parent environment
  os.chdir("/") 
  os.setsid() 
  os.umask(0) 

  # do second fork
  try: 
      pid = os.fork() 
      if pid > 0:
          # exit second parent
          sys.exit(0) 
  except OSError, e: 
      print >>sys.stderr, "fork #2 failed: %d (%s)" % (e.errno, e.strerror) 
      sys.exit(1) 

  # Make sure we don't spam the terminal
  try:
    sys.stdout.flush()
    sys.stderr.flush()
    os.dup2(os.open('/dev/null',os.O_RDONLY), sys.stdin.fileno())
    os.dup2(os.open('/dev/null',os.O_APPEND), sys.stdout.fileno())
    os.dup2(sys.stdout.fileno(), sys.stderr.fileno())
  except OSError, e:
    logging.critical("Unable to close open files: %d (%s)" % 
                     (e.errno,e.strerror))
    sys.exit(1)

  # start the daemon main loop
  logging.info("Starting %s %s" % (PROGRAM, VERSION))

def main(argv=None):
    if argv is None:
        argv = sys.argv
    parser = OptionParser(version="%prog revision "+VERSION)
    parser.add_option("", "--license", dest="show_license",
      help="show program's license and exit", action="store_true")
    parser.add_option("-f", "--file", dest="filename", 
      help="write images to FILE", metavar="FILE")
    parser.add_option("-d", "--dir", dest="directory", 
      help="write sequenced images to DIRECTORY", metavar="DIRECTORY")
    parser.add_option("-H", "--host", dest="host",
      help="get images from HOST", metavar="HOST")
    parser.add_option("-P", "--port", dest="port",
      help="get images from PORT on the host", metavar="PORT")
    parser.add_option("", "--bonjour", dest="zlabel",
      help="get images from Bonjour NAME._lapse._tcp.local.", metavar="NAME")
    parser.add_option("", "--foreground", dest="nodaemon",
      help="keep in foreground (don't daemonize)", action="store_true")
    #parser.add_option("-v", "--verbose", dest="verbose",
    #  help="generate additional debugging output", action="store_true")

    (options, args) = parser.parse_args()
    
    def error_msg(msg):
      print >> sys.stderr, msg
      parser.print_help()
      sys.exit(2)
      
    if options.show_license:
      print LICENSE
      sys.exit(0)
      
    if not options.filename:
      error_msg("*** ERROR: Must specify an output file.")
    options.filename = os.path.abspath(options.filename)
    if not os.access(os.path.dirname(options.filename), 
                     os.W_OK):
      error_msg("*** ERROR: Output file must be in a writable directory.")
    if options.directory and not os.access(os.path.abspath(options.directory),
                                           os.W_OK):
      error_msg("*** ERROR: Output directory must be a writable directory.")
    # TODO(wac): directory needs to be on the same filesystem as filename
    #            due to use of os.link() in GawkerCameraReader.ReadImage
    if not ((options.host and options.port) or options.zlabel):
      error_msg("*** ERROR: Must specify host and port or Bonjour name.")
    
    if options.zlabel and (options.host or options.port):
      error_msg("*** WARNING: Host and port are being ignored.")

    # Alright, preflight checks complete, ready to background
    if options.nodaemon:
      logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)
      logging.info("Starting %s %s (in foreground)" % (PROGRAM, VERSION))
    else:
      daemonize()
    
    host_port_info = GawkerListener(options.zlabel, options.host, options.port)
    
    # Set Zeroconf thread in motion
    if options.zlabel:
      zeroconf_server = Zeroconf.Zeroconf()
      logging.debug("Starting zeroconf service browser")
      browser = Zeroconf.ServiceBrowser(zeroconf_server, 
          "_lapse._tcp.local.", host_port_info)
    
    logging.debug("Creating reader object")
    obj = GawkerCameraReader(host_port_info, options.filename, 
                             options.directory)
    
    if obj:
      try:
        logging.debug("Starting read loop")
        obj.read_loop()
      except KeyboardInterrupt:
        return 1
      except Exception, e:
        logging.exception("Error in read loop")
        return 1
    return 0

if __name__ == "__main__": sys.exit(main())
